
from django.db.models.query import RawQuerySet
from django.db.models import sql

class MongoRawQuerySet(RawQuerySet):

    def __init__(self, model, using, *args , **kwargs ):
        self.model = model
        self.method = args[0]
        self.table = model._meta.db_table
        self._db = using

        select_sql = 'select * from %s' % self.table
        self.query = sql.RawQuery(sql=select_sql, using=self.db, params=None)

        self.raw_query = args[1:]
        self.kwargs = kwargs

        self._raw_cursor = None

        if self.method not in ['find', 'find_one']:
            self._execute()

    def __iter__(self):
        for row in self._execute():
            yield self.transform_results(row)

    def _execute(self):
        from django.db import connections
        cursor = connections[self.db].cursor()
        query = (self.table, self.method, self.raw_query, self.kwargs )

        if self._raw_cursor is None:
            self._raw_cursor = cursor.execute(query)
        return self._raw_cursor

    @property
    def columns(self):
        columns = []
        for f in self.model._meta.fields:
            columns.append( f.attname )
        return columns

    def count(self):
        return self._execute().count()
